import tushare as ts
import datetime as dt
import pandas as pd
from datetime import datetime
from django.conf import settings
import logging,time,asyncio
from asgiref.sync import sync_to_async
from .models import *
from .constants import *
from django.db.models import Max
from django.db import transaction

logger = logging.getLogger(__name__)

# 创建异步的数据库操作函数
update_data_record = sync_to_async(DataUpdateRecord.update_record)

# 定义一个同步函数来获取所有股票列表（排除ST股票）
def get_stocks_list():
    return list(StockBasic.objects.all())

# 将同步函数转换为异步函数
get_all_stocks = sync_to_async(get_stocks_list)
def get_last_update_date(data_type):
    """
    获取指定数据类型的最后更新日期
    :param data_type: 数据类型，来自DATA_TYPES常量
    :return: 最后更新日期，如果没有记录则返回START_KDATA_DATE
    """
    
    if data_type == 'index_weekly':
        start_date =  START_KDATA_DATE_WEEK
    elif data_type == 'index_monthly':
        start_date = START_KDATA_DATE_MONTH
    else:
        start_date = START_KDATA_DATE
    try:
        record = DataUpdateRecord.objects.filter(data_type=data_type).order_by('-last_update_date').first()
        return record.last_update_date if record else start_date
    except Exception as e:
        logger.error(f"Error getting last update date for {data_type}: {str(e)}")
        return start_date

# 计算两个日期之间的交易日数量
def count_trading_days(start_date, end_date):
    """
    计算两个日期之间的交易日数量
    :param start_date: 开始日期，datetime.date对象或START_KDATA_DATE常量
    :param end_date: 结束日期，datetime.date对象
    :return: 交易日数量
    """
    if not start_date or not end_date:
        return 0
    
    # 如果start_date是START_KDATA_DATE常量，直接使用它（因为它已经是datetime.date对象）
    if start_date == START_KDATA_DATE:
        start_date = START_KDATA_DATE
    
    if start_date > end_date:#起始日期和结束日期相同时可能为一天
        return 0
    
    try:
        # 将日期转换为字符串格式YYYYMMDD
        start_str = start_date.strftime('%Y%m%d')
        end_str = end_date.strftime('%Y%m%d')
        
        # 查询交易日历
        trading_days = TradeDate.objects.filter(
            cal_date__gte=start_str,
            cal_date__lte=end_str,
            is_open=True
        ).count()
        
        return trading_days
    except Exception as e:
        logger.error(f"Error counting trading days: {str(e)}")
        return 0

# 将同步函数转换为异步函数
get_last_update_date_async = sync_to_async(get_last_update_date)
count_trading_days_async = sync_to_async(count_trading_days)

class TushareAPI_Reference:
    def __init__(self, token=None):
        """
        初始化 Tushare API
        :param token: Tushare API token，如果未提供则从环境变量获取
        """
        self.token = token or getattr(settings, 'TUSHARE_TOKEN', None)
        if not self.token:
            raise ValueError("Tushare API token is required")
        
        # 初始化 Tushare API
        ts.set_token(self.token)
        self.pro = ts.pro_api()
        self.plate_to_create = []#当前更新的板块列表，作为全局变量用于更新板块的个股关系数据

    async def update_all_data(self, start_date=None, end_date=None, websocket=None):
        """
        更新所有数据
        :param start_date: 开始日期，可选
        :param end_date: 结束日期，可选
        :param websocket: WebSocket对象，可选
        """
        results = {}
        try:
            #但实际更新还是看数据库中存储的各类数据的更新日期
            logger.info(f"Starting update of all market data from {start_date} to {end_date}")
            # 转换日期为datetime.date对象，用于计算交易日
            start_date_obj = dt.datetime.strptime(start_date, '%Y-%m-%d').date() if start_date else dt.date.today()
            end_date_obj = dt.datetime.strptime(end_date, '%Y-%m-%d').date() if end_date else dt.date.today()
            
            #################################首先计算那些数据需要更新###############################################
            # 计算各数据类型的实际起始日期
            def get_actual_start_date(last_update_date):
                if not last_update_date:
                    return start_date_obj
                else:#以数据库中记录的各个表格时间为准
                    next_day = last_update_date + dt.timedelta(days=1)
                    return  next_day

            #1. 检查个股前十股东数据是否需要更新，以及流通股东数据是否需要更新
            stock_holder_top10_last_update = await get_last_update_date_async(DATA_TYPES['STOCK_HOLDER_TOP10'])
            stock_holder_top10_start_date = get_actual_start_date(stock_holder_top10_last_update)
            stock_holder_top10_trade_days = await count_trading_days_async(stock_holder_top10_start_date, end_date_obj)
            
            #2.检查大宗交易数据是否需要更新
            market_big_trade_last_update = await get_last_update_date_async(DATA_TYPES['MARKET_BIG_TRADE'])
            market_big_trade_start_date = get_actual_start_date(market_big_trade_last_update)
            market_big_trade_trade_days = await count_trading_days_async(market_big_trade_start_date, end_date_obj)
            
            #3.检查基金数据是否需要更新
            market_fund_e_last_update = await get_last_update_date_async(DATA_TYPES['MARKET_FUND_E'])
            market_fund_e_start_date = get_actual_start_date(market_fund_e_last_update)
            market_fund_e_trade_days = await count_trading_days_async(market_fund_e_start_date, end_date_obj)

            # 设置总步骤数（包括同花顺相关数据和指数数据的更新）
            total_steps = (1 if stock_holder_top10_trade_days > 0 else 0) \
                          + (1 if market_big_trade_trade_days > 0 else 0) \
                          + (1 if market_fund_e_trade_days > 0 else 0)
            # 初始化当前步骤数
            current_step = 0
            if websocket:
                await websocket.send_progress(
                    current=0,
                    total=total_steps,
                    message="开始下载数据..."
                )
            ########################################开始正式下载数据##############################################
            # 1.更新个股前十股东数据
            if websocket:
                current_step += 1
                await websocket.send_progress(
                    current=current_step,
                    total=total_steps,
                    message="正在更新个股前十及流通前十股东数据..."
                )
            print(f'更新个股前十股东数据，更新日期为{stock_holder_top10_start_date.strftime("%Y-%m-%d")}')
            if stock_holder_top10_trade_days > 0:
                results['stock_holder_top10'] = await self.update_stock_holder_top10(
                    start_date=stock_holder_top10_start_date.strftime('%Y%m%d'),
                    end_date=end_date_obj.strftime('%Y%m%d'),
                    websocket=websocket)
                logger.info("Updated stock_holder_top10 data")

                await update_data_record(DATA_TYPES['STOCK_HOLDER_TOP10'],end_date_obj)


            # 2.更新大宗交易数据
            if websocket:
                current_step += 1
                await websocket.send_progress(
                    current=current_step,
                    total=total_steps,
                    message="正在更新大宗交易数据..."
                )
            print(f'更新大宗交易数据，更新日期为{market_big_trade_start_date.strftime("%Y-%m-%d")}')
            if market_big_trade_trade_days > 0:
                results['market_big_trade'] = await self.update_market_big_trade(
                    start_date=market_big_trade_start_date.strftime('%Y%m%d'),
                    end_date=end_date_obj.strftime('%Y%m%d'),
                    websocket=websocket)
                logger.info("Updated market_big_trade data")

            # 3.更新基金数据
            if websocket:
                current_step += 1
                await websocket.send_progress(
                    current=current_step,
                    total=total_steps,
                    message="正在更新基金数据..."
                )
            print(f'更新基金数据，更新日期为{market_fund_e_start_date.strftime("%Y-%m-%d")}')
            if market_fund_e_trade_days > 0:
                results['market_fund_e'] = await self.update_market_fund_e(
                    start_date=market_fund_e_start_date.strftime('%Y%m%d'),
                    end_date=end_date_obj.strftime('%Y%m%d'),
                    websocket=websocket)
                logger.info("Updated market_fund_e data")   
                await update_data_record(DATA_TYPES['MARKET_FUND_E'],end_date_obj)
                
    

         # 更新数据更新记录（无论是否有更新都记录）
            await update_data_record(
                DATA_TYPES['ALL_DATA_4'],
                end_date_obj
            )

            # 更新数据更新记录（无论是否有更新都记录）
            await update_data_record(
                DATA_TYPES['ALL_DATA'],
                end_date_obj
            )

        except Exception as e:
            logger.error(f"Error updating all market data: {str(e)}")
            raise
                
    async def update_stock_holder_top10(self, start_date, end_date, websocket=None):
        """
        更新个股十大股东数据
        :param start_date: 开始日期(YYYYMMDD)
        :param end_date: 结束日期(YYYYMMDD) 
        :param websocket: 可选的WebSocket连接
        :return: 更新的记录数
        """
        try:
            # 计算实际查询起始日期(结束日期的2年前与开始日期的最大值)
            end_date_obj = dt.datetime.strptime(end_date, '%Y%m%d').date()
            two_years_ago = end_date_obj - dt.timedelta(days=730)  # 约2年
            start_date_obj = dt.datetime.strptime(start_date, '%Y%m%d').date()
            actual_start_date = max(two_years_ago, start_date_obj).strftime('%Y%m%d')
            
            # 获取所有股票列表
            stocks = await get_all_stocks()
            if not stocks:
                logger.warning("未获取到股票列表")
                return 0
                
            total_count = 0
            first_stock_checked = False
            
            # 定义异步数据库操作
            @sync_to_async
            def bulk_update_stock_holder(df):
                with transaction.atomic():
                    # 检查并过滤掉已存在的记录
                    existing_records = set(
                        StockHolderTop10.objects.filter(
                            ts_code__in=df['ts_code'].tolist(),
                            ann_date__in=df['ann_date'].tolist(),
                            end_date__in=df['end_date'].tolist(),
                            holder_name__in=df['holder_name'].tolist(),
                            hold_amount__in=df['hold_amount'].tolist()
                        ).values_list('ts_code', 'ann_date', 'end_date', 'holder_name','hold_amount')
                    )
                    
                    # 过滤出新记录
                    new_records = []
                    for _, row in df.iterrows():
                        key = (row['ts_code'], row['ann_date'], row['end_date'], row['holder_name'],row['hold_amount'])
                        if key not in existing_records:
                            new_records.append(row)
                    
                    if not new_records:
                        return 0
                    
                    # 只创建新记录
                    objs = [
                        StockHolderTop10(
                            ts_code=row['ts_code'],
                            ann_date=row['ann_date'],
                            end_date=row['end_date'],
                            holder_name=row['holder_name'],
                            hold_amount=row['hold_amount'],
                            hold_ratio=row['hold_ratio'],
                            hold_float_ratio=row['hold_float_ratio'],
                            hold_change=row['hold_change'],
                            holder_type=row['holder_type']
                        ) for row in new_records
                    ]
                    # objs = [
                    #     StockHolderTop10(
                    #         ts_code=row['ts_code'],
                    #         ann_date=row['ann_date'],
                    #         end_date=row['end_date'],
                    #         holder_name=row['holder_name'],
                    #         hold_amount=row['hold_amount'],
                    #         hold_ratio=row['hold_ratio'],
                    #         hold_float_ratio=row['hold_float_ratio'],
                    #         hold_change=row['hold_change'],
                    #         holder_type=row['holder_type']
                    #     ) for _, row in df.iterrows()
                    # ]
                    StockHolderTop10.objects.bulk_create(objs, batch_size=500)
                    return len(objs)
            @sync_to_async
            def bulk_update_stock_holder_float(df):
                with transaction.atomic():
                    existing_records = set(
                        StockHolderTop10Float.objects.filter(
                            ts_code__in=df['ts_code'].tolist(),
                            ann_date__in=df['ann_date'].tolist(),
                            end_date__in=df['end_date'].tolist(),
                            holder_name__in=df['holder_name'].tolist(),
                            hold_amount__in=df['hold_amount'].tolist()
                        ).values_list('ts_code', 'ann_date', 'end_date', 'holder_name','hold_amount')
                    )
                    
                    # 过滤出新记录
                    new_records = []
                    for _, row in df.iterrows():
                        key = (row['ts_code'], row['ann_date'], row['end_date'], row['holder_name'],row['hold_amount'])
                        if key not in existing_records:
                            new_records.append(row)
                    
                    if not new_records:
                        return 0
                    
                     # 只创建新记录
                    objs = [
                        StockHolderTop10Float(
                            ts_code=row['ts_code'],
                            ann_date=row['ann_date'],
                            end_date=row['end_date'],
                            holder_name=row['holder_name'],
                            hold_amount=row['hold_amount'],
                            hold_ratio=row['hold_ratio'],
                            hold_float_ratio=row['hold_float_ratio'],
                            hold_change=row['hold_change'],
                            holder_type=row['holder_type']
                        ) for row in new_records
                    ]

                    # objs = [
                    #     StockHolderTop10Float(
                    #         ts_code=row['ts_code'],
                    #         ann_date=row['ann_date'],
                    #         end_date=row['end_date'],
                    #         holder_name=row['holder_name'],
                    #         hold_amount=row['hold_amount'],
                    #         hold_ratio=row['hold_ratio'],
                    #         hold_float_ratio=row['hold_float_ratio'],
                    #         hold_change=row['hold_change'],
                    #         holder_type=row['holder_type']
                    #     ) for _, row in df.iterrows()
                    # ]
                    StockHolderTop10Float.objects.bulk_create(objs, batch_size=500)
                    return len(objs)
            
            # 遍历股票
            for stock in stocks:
                ts_code = stock.ts_code
                
                # 获取十大股东数据
                df = self.pro.top10_holders(
                    ts_code=ts_code,
                    start_date=actual_start_date,
                    end_date=end_date
                )
                df_float = self.pro.top10_floatholders(
                    ts_code=ts_code,
                    start_date=actual_start_date,
                    end_date=end_date
                )
                
                
                # 检查第一只股票是否有数据
                if not first_stock_checked:
                    if df.empty:
                        logger.info(f"首只股票{ts_code}无十大股东数据，终止更新")
                        return 0
                    first_stock_checked = True
                count = 0
                count_float = 0
                if not df.empty:
                    count = await bulk_update_stock_holder(df)
                if not df_float.empty:
                    count_float = await bulk_update_stock_holder_float(df_float)
                total_count += count+count_float
                if websocket:
                    await websocket.send_progress(
                        current=total_count,
                        message=f"更新{ts_code}的十大股东及流通十大股东数据: 新增{count}条"
                    )
            await update_data_record(
                DATA_TYPES['STOCK_HOLDER_TOP10'],
                end_date_obj
            )
            logger.info(f"成功更新{total_count}条十大股东数据")
            return total_count
            
        except Exception as e:
            logger.error(f"更新十大股东数据出错: {str(e)}")
            raise

    async def update_market_big_trade(self, start_date, end_date, websocket=None):
        """
        更新大宗交易数据
        :param start_date: 开始日期(YYYYMMDD)
        :param end_date: 结束日期(YYYYMMDD)
        :param websocket: 可选的WebSocket连接
        :return: 更新的记录数
        """
        try:
            # 解析日期
            if start_date and end_date:
                if '-' in start_date:  # 格式为YYYY-MM-DD
                    start_date_obj = dt.datetime.strptime(start_date, '%Y-%m-%d').date()
                else:  # 格式为YYYYMMDD
                    start_date_obj = dt.datetime.strptime(start_date, '%Y%m%d').date()
                
                if '-' in end_date:  # 格式为YYYY-MM-DD
                    end_date_obj = dt.datetime.strptime(end_date, '%Y-%m-%d').date()
                else:  # 格式为YYYYMMDD
                    end_date_obj = dt.datetime.strptime(end_date, '%Y%m%d').date()
            else:
                # 如果没有指定日期范围，默认更新最近10个交易日
                end_date_obj = dt.date.today()
                start_date_obj = end_date_obj - dt.timedelta(days=20)  # 多取几天以确保包含10个交易日

            # 获取交易日列表
            trade_dates = await sync_to_async(list)(TradeDate.objects.filter(
                cal_date__gte=start_date_obj.strftime('%Y%m%d'),
                cal_date__lte=end_date_obj.strftime('%Y%m%d'),
                is_open=True
            ).order_by('cal_date').values_list('cal_date', flat=True))
            
            if not trade_dates:
                logger.warning(f"没有找到{start_date}到{end_date}之间的交易日")
                return 0
                
            total_count = 0
            batch_size = 10  # 每10个交易日一批
            
            # 定义异步数据库操作
            @sync_to_async
            def bulk_update_market_big_trade(df):
                with transaction.atomic():
                    # 检查并过滤掉已存在的记录
                    existing_records = set(
                        MarketBigTrade.objects.filter(
                            trade_date__in=df['trade_date'].tolist(),
                            ts_code__in=df['ts_code'].tolist(),
                        ).values_list('trade_date', 'ts_code')
                    )
                    
                    # 过滤出新记录
                    new_records = []
                    for _, row in df.iterrows():
                        key = (row['trade_date'], row['ts_code'])
                        if key not in existing_records:
                            new_records.append(row)
                    
                    if not new_records:
                        return 0
                    
                    # 批量创建新记录
                    objs = [
                        MarketBigTrade(
                            trade_date=row['trade_date'],
                            ts_code=row['ts_code'],
                            price=row['price'],
                            vol=row['vol'],
                            amount=row['amount'],
                            buyer=row['buyer'],
                            seller=row['seller']
                        ) for row in new_records
                    ]
                    MarketBigTrade.objects.bulk_create(objs, batch_size=1000)
                    return len(objs)
            
            # 分批处理交易日
            for i in range(0, len(trade_dates), batch_size):
                batch_dates = trade_dates[i:i + batch_size]
                start_str = batch_dates[0]
                end_str = batch_dates[-1]
                
                # 获取大宗交易数据
                df = self.pro.block_trade(
                    start_date=start_str,
                    end_date=end_str
                )
                
                if not df.empty:
                    count = await bulk_update_market_big_trade(df)
                    total_count += count
                    if websocket:
                        await websocket.send_progress(
                            current=i,
                            total=len(trade_dates),
                            message=f"更新{start_str}到{end_str}的大宗交易数据: 新增{count}条"
                        )
            
            await update_data_record(
                DATA_TYPES['MARKET_BIG_TRADE'],
                end_date_obj
            )

            logger.info(f"成功更新{total_count}条大宗交易数据")
            return total_count
            
        except Exception as e:
            logger.error(f"更新大宗交易数据出错: {str(e)}")
            raise
    async def update_market_fund_e(self, start_date, end_date, websocket=None):
        """
        更新场内基金数据
        :param start_date: 开始日期(YYYYMMDD)
        :param end_date: 结束日期(YYYYMMDD)
        :param websocket: 可选的WebSocket连接
        :return: 更新的记录数
        """
        try:
            # 解析日期
            # 解析日期
            if start_date and end_date:
                if '-' in start_date:  # 格式为YYYY-MM-DD
                    start_date_obj = dt.datetime.strptime(start_date, '%Y-%m-%d').date()
                else:  # 格式为YYYYMMDD
                    start_date_obj = dt.datetime.strptime(start_date, '%Y%m%d').date()
                
                if '-' in end_date:  # 格式为YYYY-MM-DD
                    end_date_obj = dt.datetime.strptime(end_date, '%Y-%m-%d').date()
                else:  # 格式为YYYYMMDD
                    end_date_obj = dt.datetime.strptime(end_date, '%Y%m%d').date()
            else:
                # 如果没有指定日期范围，默认更新最近10个交易日
                end_date_obj = dt.date.today()
                start_date_obj = end_date_obj - dt.timedelta(days=180)  # 多取几天以确保包含10个交易日
            
            #每个月更新一次。以避免重复更新
            if end_date_obj.month == start_date_obj.month:
                return 0
            # 1. 更新基金列表
            logger.info("开始更新基金列表数据")
            if websocket:
                await websocket.send_progress(
                    current=0,
                    total=100,
                    message="正在更新基金列表..."
                )
            
            # 获取基金列表数据
            fund_list = self.pro.fund_basic(market='E', status='L')
            # 提取指定字段
            fund_list = fund_list[['ts_code','name','management','custodian','fund_type','list_date','issue_amount','status','invest_type']]
            
            # 定义异步数据库操作
            @sync_to_async
            def bulk_update_fund_list(df):
                with transaction.atomic():
                    # 获取所有现有基金记录
                    MarketFundEList.objects.all().delete()
                    # 准备批量创建和更新的对象
                    to_create = []
                    for _, row in df.iterrows():
                        row_data = row.to_dict()
                        # 创建新记录
                        to_create.append(MarketFundEList(**row_data))
                    
                    # 批量操作
                    if to_create:
                        MarketFundEList.objects.bulk_create(to_create, batch_size=500)
                    
                    return len(df)
            
            # 更新基金列表
            fund_count = await bulk_update_fund_list(fund_list)
            logger.info(f"成功更新{fund_count}条基金列表数据")
            
            # 2. 更新基金持仓数据
            logger.info("开始更新基金持仓数据")
            if websocket:
                await websocket.send_progress(
                    current=50,
                    total=100,
                    message="正在更新基金持仓数据..."
                )
            
            # 定义异步数据库操作
            @sync_to_async
            def bulk_update_fund_detail(df,ts_code):
                with transaction.atomic():
                    # 检查并过滤掉已存在的记录
                    MarketFundEDetail.objects.filter(ts_code=ts_code).delete()

                    to_create = []
                    for _, row in df.iterrows():
                        row_data = row.to_dict()
                        # 创建新记录
                        to_create.append(MarketFundEDetail(**row_data))
                    
                    # 批量操作
                    if to_create:
                        MarketFundEDetail.objects.bulk_create(to_create, batch_size=500)
                    return len(df)
            
            # 遍历所有基金代码获取持仓数据
            total_detail_count = 0
            all_fund_codes = fund_list['ts_code'].tolist()
            
            for i, ts_code in enumerate(all_fund_codes):
                print(f'正在更新{ts_code}的持仓数据 end_date = {end_date}')
                try:
                    # 获取基金持仓数据
                    df = self.pro.fund_portfolio(
                        ts_code=ts_code,
                        end_date=end_date
                    )
                    print(f'df = {df}')
                    
                    if not df.empty:
                        count = await bulk_update_fund_detail(df,ts_code)
                        total_detail_count += count
                        
                        if websocket:
                            progress = 50 + int(50 * (i + 1) / len(all_fund_codes))
                            await websocket.send_progress(
                                current=progress,
                                total=100,
                                message=f"更新{ts_code}的持仓数据: 新增{count}条"
                            )
                    else:
                        logger.warning(f"没有找到{ts_code}的持仓数据")
                        continue
                except Exception as e:
                    logger.error(f"更新基金{ts_code}持仓数据出错: {str(e)}")
            

            await update_data_record(
                DATA_TYPES['MARKET_FUND_E'],
                end_date_obj
            )

            logger.info(f"成功更新{total_detail_count}条基金持仓数据")
            return {"fund_list": fund_count, "fund_detail": total_detail_count}
            
        except Exception as e:
            logger.error(f"更新基金数据出错: {str(e)}")
            raise       